{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.stats as stats\n",
    "import cvxpy as cvx\n",
    "import plotly as py\n",
    "import plotly.graph_objs as go\n",
    "import helper\n",
    "\n",
    "py.offline.init_notebook_mode(connected=True)\n",
    "%matplotlib inline\n",
    "plt.style.use('ggplot')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization with an Alpha Model and a Risk Model\n",
    "In this exercise, we will construct an optimization problem using a stock universe consisting of 3 stocks. This problem will inherently be somewhat artificial, but using a stock universe of 3 stocks will allow us to create plots in 3 dimensions. These are useful for illustrating exactly what is going on in each step and will help develop your intuition for the problem. In this exercise, we will do the following steps:\n",
    "1. Create a 1-year-momentum-based alpha vector from some stock price data. \n",
    "2. Create a risk model using PCA.\n",
    "3. Construct the optimization problem using the alpha vector and risk model, and apply a set of standard constraints.\n",
    "4. Solve the optimization problem.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data\n",
    "Load the data from the file `stock_prices_advanced_optimization.csv`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prices = pd.read_csv('stock_prices_advanced_optimization.csv', parse_dates=['date'], index_col=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data are the price trends for 3 stocks, Stock A, Stock B and Stock C, for 4 years, from 2013 to 2017."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "prices.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "prices.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the returns of these price data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rets = prices.pct_change()[1:].fillna(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an alpha vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from scipy.stats import zscore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll create an alpha vector based on 1-year momentum. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 1-yr momentum alpha\n",
    "\n",
    "def log_returns(series):\n",
    "    return np.log(series[-1])-np.log(series[0])\n",
    "    \n",
    "alpha = prices.rolling(window=252).apply(log_returns).rank(axis='columns').apply(zscore, axis='columns')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll take the row of most recent alpha values to use as our alpha vector. There should be a value in the vector for each stock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Take most recent set of values\n",
    "alpha_vector = alpha.iloc[-1,:]\n",
    "print(alpha_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this problem, we'll use the optimization objective $-\\boldsymbol{\\alpha}^T \\mathbf{x}$. Remember, we are trying to _minimize_ this function (to maximize alpha). Let's plot $-\\boldsymbol{\\alpha}^T \\mathbf{x}$ as a function of the components of $\\mathbf{x}$ so that we can get a better sense of the \"shape\" of the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n = 10\n",
    "x = y = z = np.linspace(-2,2,n)\n",
    "\n",
    "xv, yv, zv = np.meshgrid(x, y, z, indexing='ij')\n",
    "obj_val = np.full(xv.shape, np.nan)\n",
    "\n",
    "for i in range(n):\n",
    "    for j in range(n):\n",
    "        for k in range(n):\n",
    "            obj_val[i,j,k] = -alpha_vector[0]*xv[i,j,k]-alpha_vector[1]*yv[i,j,k]-alpha_vector[2]*zv[i,j,k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "hover_text = helper.generate_hover_text(xv, yv, zv, 'Weight on Stock A', 'Weight on Stock B', 'Weight on Stock C', fcn_values=obj_val, fcn_label='Objective Function')\n",
    "\n",
    "trace1 = go.Scatter3d(\n",
    "    x=xv.flatten(),\n",
    "    y=yv.flatten(),\n",
    "    z=zv.flatten(),\n",
    "    text = hover_text.flatten(),\n",
    "    mode='markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color=obj_val.flatten(),     # set color to an array/list of desired values\n",
    "        colorscale='Viridis',   # choose a colorscale\n",
    "        colorbar=dict(\n",
    "                title='Objective Function'\n",
    "            ),\n",
    "        opacity=0.4,\n",
    "        showscale=True\n",
    "    ),\n",
    "    hoverinfo = 'text'\n",
    ")\n",
    "\n",
    "data = [trace1]\n",
    "layout = helper.create_standard_layout('Alpha Vector Optimization Objective Function', 'Weight on')\n",
    "fig = go.Figure(data=data, layout=layout)\n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this plot, the value of the objective function is represented by the color at each point in weight space. Notice that the function increases with weight on Stock C, decreases with weight on Stock B, and doesn't depend on the weight on Stock A."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens if we try to run the optimization now, by trying to minimize this objective function, under no constraints?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def find_optimal_holdings(alpha_vector):\n",
    "    \"\"\"\n",
    "    Use cvxpy to construct and solve an optimization problem. Use -alpha*x as the objective, and use no constraints.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    alpha_vector : DataFrame\n",
    "        The 3-stock alpha vector calculated above.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    optimal_weights : DataFrame\n",
    "        A DataFrame containing the optimal weights calculated by the optimization algorithm.\n",
    "    \"\"\"\n",
    "    #TODO: Implement function\n",
    "    \n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimal_weights = find_optimal_holdings(\n",
    "    alpha_vector\n",
    ")\n",
    "print(\"Optimal weights: \", optimal_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, without constraints, the problem is unbounded. We could achieve a smaller and smaller objective function value by taking larger and larger long positions on Stock C, and larger and larger short positions on Stock B. But this would increase risk and leverage. This is where our constraints come into play."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Risk Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to create a PCA-based risk model for our 3 stocks. First, let's plot the returns data and their mean, just to have a sense for the shape of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "m = rets.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Trace for mean return vector m\n",
    "mean_vec = np.vstack((np.full(3, 0), m.values)).T\n",
    "\n",
    "hover_text2 = helper.generate_hover_text(mean_vec[0], mean_vec[1], mean_vec[2], 'Mean of Returns of Stock A', 'Mean of Returns of Stock B', 'Mean of Returns of Stock C', sig_figs = 7)\n",
    "\n",
    "trace2 = go.Scatter3d(\n",
    "    x=mean_vec[0],\n",
    "    y=mean_vec[1],\n",
    "    z=mean_vec[2],\n",
    "    mode='lines+markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color='red',\n",
    "        opacity=0.9,\n",
    "        symbol=\"diamond\"\n",
    "    ),\n",
    "    name = 'mean daily return',\n",
    "    text = hover_text2.flatten(),\n",
    "    hoverinfo = 'text'\n",
    ")\n",
    "\n",
    "\n",
    "# Trace for data\n",
    "hover_text3 = helper.generate_hover_text(rets['A'].values, rets['B'].values, rets['C'].values, 'Return of Stock A', 'Return of Stock B', 'Return of Stock C')\n",
    "\n",
    "trace3 = go.Scatter3d(\n",
    "    x=rets['A'].values,\n",
    "    y=rets['B'].values,\n",
    "    z=rets['C'].values,\n",
    "    mode='markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color='#7FB3D5',  \n",
    "        opacity=0.3,\n",
    "    ),\n",
    "    name = 'daily returns',\n",
    "    text = hover_text3.flatten(),\n",
    "    hoverinfo = 'text'\n",
    ")\n",
    "\n",
    "data = [trace2, trace3]\n",
    "\n",
    "layout = helper.create_standard_layout('Returns Data', 'Return of')\n",
    "fig = go.Figure(data=data, layout=layout)\n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a few things to notice here. \n",
    "1. The vector of mean returns is very small compared to the volatility of the data! Admittedly, this is an artificial example with fake data, but this is also generally the case.\n",
    "2. The volatility of Stock C is larger than the volatility of Stock B, which is larger than the volatility of Stock A."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given these data, we can already set up some expectations for what the results of PCA will look like. Remember that we expect the first PC to point in the direction of maximum variability of the data. Given the volatility of Stock C, we expect the first PC to point approximately parallel to the \"Stock C axis\". If we get results that are very different from this, we know that something is wrong with what we are doing. Let's now run the PCA algorithm, and see what we get. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we'll mean-center the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rets = rets - m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we'll create a class that will fit our risk model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "class RiskModelPCA():\n",
    "    \n",
    "    ANN_FACTOR = 252\n",
    "    \n",
    "    def __init__(self, num_factors):\n",
    "        self._num_factors = num_factors\n",
    "        self.num_stocks_ = None\n",
    "        self.factor_betas_ = None\n",
    "        self.factor_returns_ = None\n",
    "        self.common_returns_ = None\n",
    "        self.residuals_ = None\n",
    "        self.factor_cov_matrix_ = None\n",
    "        self.idio_var_matrix_ = None\n",
    "        self.explained_variance_ratio_ = None\n",
    "\n",
    "    def fit(self, returns):\n",
    "        self.num_stocks_ = len(returns.columns)\n",
    "        mod = PCA(n_components=self._num_factors, svd_solver='full')\n",
    "        mod.fit(returns)\n",
    "        \n",
    "        self.factor_betas_ = pd.DataFrame(\n",
    "            data=mod.components_.T,\n",
    "            index=returns.columns\n",
    "        )\n",
    "        \n",
    "        self.factor_returns_ = pd.DataFrame(\n",
    "            data=mod.transform(rets),\n",
    "            index=returns.index\n",
    "        )\n",
    "        \n",
    "        self.explained_variance_ratio_ = mod.explained_variance_ratio_\n",
    "        \n",
    "        self.common_returns_ = pd.DataFrame(\n",
    "            data=np.dot(self.factor_returns_, self.factor_betas_.T),\n",
    "            index=returns.index\n",
    "        )\n",
    "        self.common_returns_.columns = returns.columns\n",
    "        \n",
    "        self.residuals_ = (returns - self.common_returns_)\n",
    "        \n",
    "        self.factor_cov_matrix_ = np.diag(\n",
    "            self.factor_returns_.var(axis=0, ddof=1)*RiskModelPCA.ANN_FACTOR\n",
    "        )\n",
    "        \n",
    "        self.idio_var_matrix_ = pd.DataFrame(\n",
    "            data=np.diag(np.var(self.residuals_))*RiskModelPCA.ANN_FACTOR,\n",
    "            index=returns.columns\n",
    "        )\n",
    "        \n",
    "        self.idio_var_vector_ = pd.DataFrame(\n",
    "            data=np.diag(self.idio_var_matrix_.values),\n",
    "            index=returns.columns\n",
    "        )\n",
    "        \n",
    "        self.idio_var_matrix_.columns = index=returns.columns\n",
    "\n",
    "    def get_factor_exposures(self, weights):\n",
    "        B = self.factor_betas_.loc[weights.index]\n",
    "        return B.T.dot(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's fit the risk model with 2 factors (i.e., we'll keep 2 PCs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rm = RiskModelPCA(2) # create an instance of the class with 2 factors\n",
    "rm.fit(rets) # fit the model on our data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the factors (PCs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rm.factor_betas_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first PC points almost completely in the \"Stock C\" direction. The second points almost completely in the \"Stock B\" direction. This makes sense given the fact that we see most variability in the \"Stock C\" direction, and if we collapse that variability, we'd see the next most variability in the \"Stock B\" direction. Let's plot the PC vectors to make this more clear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PC_scaler = 0.04 # The PC vectors have length 1, but this is larger than the scale of our data, so for visualization purposes, let's plot scaled-down versions of the PCs. \n",
    "\n",
    "# Trace for PC 0\n",
    "pc0 = np.vstack((np.full(3, 0), rm.factor_betas_[0].values)).T*PC_scaler\n",
    "\n",
    "hover_text4 = helper.generate_hover_text(pc0[0], pc0[1], pc0[2], 'Return of Stock A', 'Return of Stock B', 'Return of Stock C')\n",
    "\n",
    "trace4 = go.Scatter3d(\n",
    "    x=pc0[0],\n",
    "    y=pc0[1],\n",
    "    z=pc0[2],\n",
    "    mode='lines+markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color='#45B39D',\n",
    "        opacity=0.9,\n",
    "        symbol=\"diamond\"\n",
    "    ),\n",
    "    line=dict(\n",
    "        color='#45B39D',\n",
    "        width=3\n",
    "    ),\n",
    "    name = 'PC 0',\n",
    "    text = hover_text4.flatten(),\n",
    "    hoverinfo = 'text'\n",
    "\n",
    ")\n",
    "\n",
    "# Trace for PC 1\n",
    "pc1 = np.vstack((np.full(3, 0), rm.factor_betas_[1].values)).T*PC_scaler\n",
    "\n",
    "hover_text5 = helper.generate_hover_text(pc1[0], pc1[1], pc1[2], 'Return of Stock A', 'Return of Stock B', 'Return of Stock C')\n",
    "\n",
    "trace5 = go.Scatter3d(\n",
    "    x=pc1[0],\n",
    "    y=pc1[1],\n",
    "    z=pc1[2],\n",
    "    mode='lines+markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color='#FFC300',\n",
    "        opacity=0.9,\n",
    "        symbol=\"diamond\"\n",
    "    ),\n",
    "    line=dict(\n",
    "        color='#FFC300',\n",
    "        width=3\n",
    "    ),\n",
    "    name = 'PC 1',\n",
    "    text = hover_text5.flatten(),\n",
    "    hoverinfo = 'text'\n",
    "\n",
    ")\n",
    "\n",
    "# Trace for data\n",
    "hover_text6 = helper.generate_hover_text(rets['A'].values, rets['B'].values, rets['C'].values, 'Return of Stock A', 'Return of Stock B', 'Return of Stock C')\n",
    "\n",
    "trace6 = go.Scatter3d(\n",
    "    x=rets['A'].values,\n",
    "    y=rets['B'].values,\n",
    "    z=rets['C'].values,\n",
    "    mode='markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color='#7FB3D5', \n",
    "        opacity=0.3,\n",
    "    ),\n",
    "    name = 'daily returns',\n",
    "    text = hover_text6.flatten(),\n",
    "    hoverinfo = 'text'\n",
    ")\n",
    "\n",
    "data = [trace4, trace5, trace6]\n",
    "\n",
    "layout = helper.create_standard_layout('Returns Data with Factor (PC) Directions', 'Return of')\n",
    "fig = go.Figure(data=data, layout=layout)\n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the fraction of variance explained by the two factors we kept. This should confirm our intuition about the relative scale of variance explained by the first two factors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.bar(np.arange(2), rm.explained_variance_ratio_, color=['#45B39D', '#FFC300']);\n",
    "plt.title('% of Variance Explained by Each Factor');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at the factor returns. We'll convert them to \"price series\" and visualize the evolution of these over time. Recall that these are the data expressed in the \"factor\" (PC) basis––the projections of the data vectors onto the factor dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(rm.factor_returns_ + 1).cumprod().plot(color=['#45B39D', '#FFC300'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in this example, since the first factor is so close to the direction of the return on Stock C, the factor return (converted to a \"price series\") for factor 0 looks a lot like the price evolution of Stock C, but _inverted_. This is because the PCs are vectors that represent _directions_ in space—they are only unique up to a sign change, and thus the projections onto them are unique up to a sign change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prices.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Optimization Constraints\n",
    "As we saw above, we need to constrain our optimization problem. Let's now create and plot the following constraints:\n",
    "* risk, based on our risk model\n",
    "* leverage\n",
    "* market neutral\n",
    "* factor exposure limits\n",
    "* individual weight limits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Risk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "B = rm.factor_betas_\n",
    "F = rm.factor_cov_matrix_\n",
    "S = np.diag(rm.idio_var_vector_.values.flatten())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the $\\mathbf{B}$, $\\mathbf{F}$, and $\\mathbf{S}$ matrices, let's write a function that takes in the portfolio weights and spits out the portfolio risk. Remember that the portfolio risk is calculated $\\mathbf{x}^T(\\mathbf{B}^T\\mathbf{F}\\mathbf{B} + \\mathbf{S})\\mathbf{x}$, but be careful to ensure the matrices are all oriented correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def risk_fcn(x):\n",
    "    \"\"\"\n",
    "    Calculate portfolio risk.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : numpy array\n",
    "        The vector of portfolio weights.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    risk : float\n",
    "        Portfolio risk.\n",
    "    \"\"\"\n",
    "    #TODO: Implement function\n",
    "    \n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_samples = 25\n",
    "grid_max = 2.5\n",
    "risk_grid, spacing, xv, yv, zv = helper.evaluate_fcn_on_grid(grid_max, n_samples, risk_fcn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at a plot of the value of the portfolio risk as a function of the portfolio weights in 3 dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hover_text = helper.generate_hover_text(xv, yv, zv, 'Weight on Stock A', 'Weight on Stock B', 'Weight on Stock C', fcn_values=risk_grid, fcn_label='Portfolio Risk')\n",
    "\n",
    "trace7 = go.Scatter3d(\n",
    "    x=xv.flatten(),\n",
    "    y=yv.flatten(),\n",
    "    z=zv.flatten(),\n",
    "    mode='markers',\n",
    "    marker=dict(\n",
    "        size=4,\n",
    "        color=risk_grid.flatten(), \n",
    "        colorscale='Viridis',\n",
    "        opacity=0.3,\n",
    "        showscale=True\n",
    "    ),\n",
    "    text = hover_text.flatten(),\n",
    "    hoverinfo = 'text'\n",
    ")\n",
    "\n",
    "data = [trace7]\n",
    "layout = helper.create_standard_layout('Portfolio Risk as Modeled by our PCA Risk Model', 'Weight on')\n",
    "fig = go.Figure(data=data, layout=layout)\n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This plot uses a color at several points in 3-D space to represent the portfolio risk as a function of weight on each stock. Note that it increases with weight on each stock. This helps us visualize the risk function, but what does it look like to _constrain_ risk? Let's visualize the boundary of the region within which risk is below a limiting value. This will help us visualize the shape of the space we search with optimization. We will do this for each of our constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "risk_data = helper.plot_iso_surface(risk_grid, 0.05, 25, 2.5, \"Portfolio Risk = 0.05\", '#F5B7B1', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Leverage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's write a function that takes in the portfolio weights and spits out the portfolio leverage, $\\sum_i|x_i|$. This will allow us to visualize the space that satisfies the constraint $\\sum_i|x_i| \\leq 1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lev_fcn(x):\n",
    "    \"\"\"\n",
    "    Calculate portfolio leverage.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : numpy array\n",
    "        The vector of portfolio weights.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    leverage : float\n",
    "        Portfolio risk.\n",
    "    \"\"\"\n",
    "    #TODO: Implement function\n",
    "    \n",
    "    return None\n",
    "\n",
    "n_samples = 25\n",
    "grid_max = 2.5\n",
    "\n",
    "lev_grid, _, _, _, _ = helper.evaluate_fcn_on_grid(grid_max, n_samples, lev_fcn)\n",
    "lev_data = helper.plot_iso_surface(lev_grid, 1.0, n_samples, grid_max, \"Leverage Ratio = 1\", '#1F618D', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Market Neutral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the market neutral constraint, we constrain the sum of the weights. Let's write a function that calculates this quantity, given the portfolio weights. Then we can visualize the plane $\\sum_i x_i = 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mn_fcn(x):\n",
    "    \"\"\"\n",
    "    Calculate the sum of the portfolio weights.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : numpy array\n",
    "        The vector of portfolio weights.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    sum_of_weights : float\n",
    "        Sum of the portfolio weights.\n",
    "    \"\"\"\n",
    "    #TODO: Implement function\n",
    "    \n",
    "    return None\n",
    "\n",
    "n_samples = 25\n",
    "grid_max = 2.5\n",
    "\n",
    "mn_grid, _, _, _, _ = helper.evaluate_fcn_on_grid(grid_max, n_samples, mn_fcn)\n",
    "mn_data = helper.plot_iso_surface(mn_grid, 0, n_samples, grid_max, \"Sum of Weights = 0\", '#D35400', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Factor Exposures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also calculate the factor exposures. Then we'll plot the planes defined by the factor exposure limits. The optimization will constrain the solution to lie between each pair of planes. We'll constrain each factor exposure to be between -0.1 and 0.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fac_fcn(x):\n",
    "    \"\"\"\n",
    "    Calculate portfolio factor exposures.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : numpy array\n",
    "        The vector of portfolio weights.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    fac_exposures : numpy array\n",
    "        Portfolio factor exposures.\n",
    "    \"\"\"\n",
    "    #TODO: Implement function\n",
    "    \n",
    "    return None\n",
    "\n",
    "grid_max = 2.5\n",
    "n_samples = 25\n",
    "\n",
    "fac_grid, spacing, _, _, _ = helper.evaluate_fcn_on_grid(grid_max, n_samples, fac_fcn)\n",
    "\n",
    "factor_limit_data = []\n",
    "\n",
    "for factor in range(B.shape[1]):\n",
    "    for l in range(2):\n",
    "        mult_fac = l*2-1\n",
    "        factor_limit_data.extend(helper.plot_iso_surface(fac_grid[:,:,:,factor], 0.1*mult_fac, n_samples, grid_max, 'Factor Exposure Limits', '#D2B4DE', False))\n",
    "\n",
    "layout = helper.create_standard_layout('Factor Exposure Limits', 'Weight on')\n",
    "\n",
    "fig = go.Figure(data=factor_limit_data, layout=layout)         \n",
    "        \n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Individual Weight Limits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's look at the space inside the constraints on each individual weight. We'll constrain each individual weight to be between -0.55 and 0.55."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_max = 0.55\n",
    "\n",
    "x = np.array([-1, -1, 1, 1, -1, -1, 1, 1])*x_max\n",
    "y = np.array([-1, 1, 1, -1, -1, 1, 1, -1])*x_max\n",
    "z = np.array([-1, -1, -1, -1, 1, 1, 1, 1])*x_max\n",
    "hover_text = helper.generate_hover_text(x, y, z, 'Weight on Stock A', 'Weight on Stock B', 'Weight on Stock C')\n",
    "\n",
    "\n",
    "weight_data = [\n",
    "    go.Mesh3d(\n",
    "        x = x,\n",
    "        y = y,\n",
    "        z = z,\n",
    "        colorscale = '#FCF3CF',\n",
    "        intensity = np.full(8, 1),\n",
    "        i = [7, 0, 0, 0, 4, 4, 6, 6, 4, 0, 3, 2],\n",
    "        j = [3, 4, 1, 2, 5, 6, 5, 2, 0, 1, 6, 3],\n",
    "        k = [0, 7, 2, 3, 6, 7, 1, 1, 5, 5, 7, 6],\n",
    "        name='Weight on Stock B',\n",
    "        showscale=False,\n",
    "        opacity = 0.3,\n",
    "        hoverinfo='none'\n",
    "    )\n",
    "]\n",
    "\n",
    "trace = go.Scatter3d(\n",
    "    x=x,\n",
    "    y=y,\n",
    "    z=z,\n",
    "    mode='markers',\n",
    "    marker=dict(\n",
    "        size=6,\n",
    "        opacity=0.0001,\n",
    "        color='#BFB1A8', \n",
    "    ),\n",
    "    text = hover_text.flatten(),\n",
    "    hoverinfo = 'text',\n",
    "    showlegend=False\n",
    ")\n",
    "    \n",
    "weight_data = [weight_data[0], trace]    \n",
    "    \n",
    "layout = helper.create_standard_layout('Individual Weight Limits', 'Weight on')\n",
    "fig = go.Figure(data=weight_data, layout=layout)\n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot all the constraints on the same axes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's visualize the space that satisfies all the constraints. It has to be inside all of the enclosed spaces, _on_ the plane defined by the market neutral constraint, and _between_ the planes defined by the factor exposure constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = risk_data\n",
    "data.extend(lev_data)\n",
    "data.extend(mn_data)\n",
    "data.extend(factor_limit_data)\n",
    "data.extend(weight_data)\n",
    "layout = helper.create_standard_layout('Visualize Intersection of All Constraints', 'Weight on')\n",
    "fig = go.Figure(data=data, layout=layout)\n",
    "py.offline.iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization\n",
    "Finally, let's define and run the optimization problem with the same objective function as above, but this time, with all the constraints we just discussed. Remember that the optimization is looking for the point in the space that satisfies all the constraints that _minimizes_ the alpha function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def find_optimal_holdings(\n",
    "    alpha_vector,\n",
    "    risk_model,\n",
    "    risk_cap=0.05,\n",
    "    factor_max=10.0,\n",
    "    factor_min=-10.0,\n",
    "    x_max=0.55,\n",
    "    x_min=-0.55):\n",
    "    \n",
    "    \"\"\"\n",
    "    Define an optimization problem. It takes in several inputs and optimization parameters and outputs the\n",
    "    optimized weights. The objective should minimize the objective -alpha*x, but also apply the risk, leverage,\n",
    "    market neutral, factor exposure, and invidiual weight constraints.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    alpha_vector : numpy array\n",
    "        The alpha vector used in the optimization objective.\n",
    "    risk_model : RiskModelPCA\n",
    "        The risk model calculated above using PCA.\n",
    "    risk_cap : float\n",
    "        The limit to be placed on portfolio risk.\n",
    "    factor_max : float\n",
    "        The factor exposure upper limit.\n",
    "    factor_min : float\n",
    "        The factor exposure lower limit.\n",
    "    x_max : float\n",
    "        The individual weight upper limit.\n",
    "    x_min : float\n",
    "        The individual weight lower limit.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    optimal_weights : numpy array\n",
    "        The optimal portfolio weights.\n",
    "    \"\"\"\n",
    "    #TODO: Implement function\n",
    "    \n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimal_weights = find_optimal_holdings(\n",
    "    alpha_vector,\n",
    "    rm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimal weights are 0 weight on Stock A, 50% long on Stock B and 50% short on Stock C."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "The solution for this notebook is [here](Advanced_Opt_solution.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
